package org.jflame.commons.crypto;

import java.security.Key;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEParameterSpec;

/**
 * 加解密需要的参数
 * 
 * @author charles.zhang
 */
public abstract class CipherArgs {

    protected AlgorithmParameterSpec initParams;

    /**
     * 返回算法初始时的附带参数
     * 
     * @see Cipher#init(int, Key, AlgorithmParameterSpec)
     * @return
     */
    public AlgorithmParameterSpec getInitParams() {
        return initParams;
    }

    public void setInitParams(AlgorithmParameterSpec initParams) {
        this.initParams = initParams;
    }

    /**
     * 设置分组需要的向量
     */
    public void setIv(byte[] iv) {
        initParams = new IvParameterSpec(iv);
    }

    /**
     * 返回密钥
     * 
     * @param cipher_mode Cipher.ENCRYPT_MODE or Cipher.DECRYPT_MODE
     * @return Cipher.ENCRYPT_MODE返回加密密钥,Cipher.DECRYPT_MODE返回解密密钥
     */
    public abstract Key getKey(int cipher_mode);

    public static SymmetricArgs buildSym(final Alg alg, String secretKeyStr) throws EncryptException {
        return new SymmetricArgs(alg, secretKeyStr);
    }

    public static SymmetricArgs buildSym(final Alg alg, String secretKeyStr, AlgorithmParameterSpec initParams)
            throws EncryptException {
        return new SymmetricArgs(alg, secretKeyStr).initParam(initParams);
    }

    public static SymmetricArgs buildSym(final Alg alg, byte[] secretRawKey) throws EncryptException {
        return new SymmetricArgs(alg, secretRawKey);
    }

    public static AsymmetricArgs buildAsym(final Alg alg, final String publicKeyStr,
            final String privateKeyStr) throws EncryptException {
        return new AsymmetricArgs(alg, publicKeyStr, privateKeyStr);
    }

    /**
     * 对称加密算法需要的参数
     * 
     * @author charles.zhang
     */
    public static class SymmetricArgs extends CipherArgs {

        private SecretKey secretKey;// 加密key

        public SymmetricArgs(final SecretKey secretKey) {
            this.secretKey = secretKey;
        }

        public SymmetricArgs(final Alg alg, byte[] secretRawKey) throws EncryptException {
            secretKey = KeyUtils.generate(alg, secretRawKey);
        }

        /**
         * 构造函数
         * 
         * @param alg 加密算法
         * @param secretKeyStr 密钥
         * @throws EncryptException 构造SecretKey对象失败抛出异常
         */
        public SymmetricArgs(final Alg alg, String secretKeyStr) throws EncryptException {
            secretKey = KeyUtils.generate(alg, secretKeyStr);
        }

        @Override
        public Key getKey(int cipher_mode) {
            return secretKey;
        }

        public SymmetricArgs iv(byte[] iv) {
            setIv(iv);
            return this;
        }

        /**
         * 设置PBE算法初始参数
         * 
         * @param salt 随机盐,长度8
         * @param iterationCount 迭代次数
         */
        public SymmetricArgs PBEInitParam(final byte[] salt, final int iterationCount) {
            super.initParams = new PBEParameterSpec(salt, iterationCount);
            return this;
        }

        /**
         * 设置PBES2算法初始参数
         * 
         * @param salt 随机盐,长度8
         * @param iterationCount 迭代次数
         * @param iv 加密向量,使用aes需要
         */
        public SymmetricArgs PBEInitParam(final byte[] salt, final int iterationCount, final byte[] iv) {
            super.initParams = new PBEParameterSpec(salt, iterationCount, new IvParameterSpec(iv));
            return this;
        }

        public SymmetricArgs initParam(AlgorithmParameterSpec algParamSpec) {
            super.initParams = algParamSpec;
            return this;
        }
    }

    /**
     * 非对称算法需要的参数
     * 
     * @author charles.zhang
     */
    public static class AsymmetricArgs extends CipherArgs {

        private PublicKey publicKey;
        private PrivateKey privateKey;

        public AsymmetricArgs(final KeyPair keyPair) {
            publicKey = keyPair.getPublic();
            privateKey = keyPair.getPrivate();
        }

        public AsymmetricArgs(final PublicKey _publicKey, final PrivateKey _privateKey) {
            if (_publicKey == null && _privateKey == null) {
                throw new IllegalArgumentException("publicKey and privateKey both cannot be empty");
            }
            this.publicKey = _publicKey;
            this.privateKey = _privateKey;
        }

        /**
         * 构造函数
         * 
         * @param alg 非对称算法
         * @param publicKeyStr 公钥文件路径(以file:或classpath:开头)或公钥base64,只用于解密时可以传null
         * @param privateKeyStr 私钥文件路径(以file:或classpath:开头)或私钥base64, 只用于加密时可以传null
         * @throws EncryptException
         */
        public AsymmetricArgs(final Alg alg, final String publicKeyStr, final String privateKeyStr)
                throws EncryptException {
            if (publicKeyStr == null && privateKeyStr == null) {
                throw new IllegalArgumentException("publicKey and privateKey both cannot be empty");
            }
            if (publicKeyStr != null) {
                this.publicKey = KeyUtils.loadPublicKey(alg, publicKeyStr);
            }
            if (privateKeyStr != null) {
                this.privateKey = KeyUtils.loadPrivateKey(alg, privateKeyStr);
            }

        }

        @Override
        public Key getKey(int cipher_mode) {
            return cipher_mode == Cipher.ENCRYPT_MODE ? publicKey : privateKey;
        }

        public AsymmetricArgs iv(byte[] iv) {
            setIv(iv);
            return this;
        }

        public AsymmetricArgs initParams(AlgorithmParameterSpec initParams) {
            super.setInitParams(initParams);
            return this;
        }

        public PublicKey getPublicKey() {
            return publicKey;
        }

        public PrivateKey getPrivateKey() {
            return privateKey;
        }

    }

}
